#include <linux/module.h>	// Needed by all modules
#include <linux/kernel.h>	// Needed for KERN_INFO
#include <linux/file.h>

#ifdef __KERNEL__
#include <linux/fs.h>
#include <linux/binfmts.h>
#else
#include "/usr/src/linux/include/linux/fs.h"
#include "/usr/src/linux/include/linux/binfmts.h"
#endif

#include "binfmt_pe.h"

static int load_pe_binary(struct linux_binprm * bprm, struct pt_regs * regs);
static int load_pe_library(struct file *file);
static int pe_core_dump(long signr, struct pt_regs *regs, struct file *file, unsigned long limit);

static struct linux_binfmt pe_format = {
    .module 	    = THIS_MODULE,
    .load_binary	= load_pe_binary,
    .load_shlib	  = load_pe_library,
    .core_dump	  = pe_core_dump,
    .min_coredump	= PE_EXEC_PAGESIZE,
    .hasvdso	    = 1
};

const char* hex = "0123456789ABCDEF"; // 1 byte wasted for null terminator but it's much cleaner

// int kernel_read(struct file *source, uint32_t begin, char *destination, uint32_t length);

#define BAILOUT_ERR(condition, err_val, goto_name) if(condition) { printk("\nbinfmt_pe loader failure\nfile: %s\nline: %i\ncondition: %s\nerror: %s", __FILE__, __LINE__, #condition, #err_val); retval = -err_val; goto goto_name; }
#define BAILOUT(condition) BAILOUT_ERR(condition, ENOEXEC, out)
#define BAILOUT_TO(condition, goto_name) BAILOUT_ERR(condition, ENOEXEC, goto_name)
#define BAILOUT_MEM(mem) BAILOUT_ERR( mem == NULL, ENOMEM, free_mem_out)
#define BAILOUT_READ(retval) BAILOUT_ERR( retval == 0, ENOEXEC, free_mem_out)

static int load_pe_binary(struct linux_binprm* bprm, struct pt_regs* regs)
{
  char* mem = NULL;
  int retval = 0;
  uint16_t i;
  size_t mem_size;
  struct s_dos_header* dos_header;
  struct s_nt_header* nt_header = NULL;
  struct s_section_header* section_headers = NULL;

// basic DOS type executable check
  dos_header = (struct s_dos_header *) bprm->buf;
  BAILOUT(dos_header->Signature != DOS_SIGNATURE)

  printk("\nnt header address: %x", dos_header->NtHeaderAddress);

// allocate mem for nt_header
  mem = kmalloc(sizeof(struct s_nt_header), GFP_KERNEL);
  BAILOUT_MEM(mem)
  nt_header = (struct s_nt_header*)mem;

// read the full NT header
  retval = kernel_read(bprm->file, dos_header->NtHeaderAddress, mem, sizeof(struct s_nt_header));
  BAILOUT_READ(retval)

// basic PE type header check
  BAILOUT(nt_header->Signature != PE_SIGNATURE)

// check if it's an executable
  BAILOUT(!(nt_header->FileHeader.Characteristics & EXECUTABLE_IMAGE))

// if this fails, the executable file is corrupt
  BAILOUT(nt_header->FileHeader.SizeOfOptionalHeader != sizeof(struct s_optional_header))

// check for matching architechture
  BAILOUT(nt_header->FileHeader.ArchitectureId != X86_ARCH)

  mem_size = nt_header->OptionalHeader.SizeOfImage;

  kfree(mem);
  mem = kmalloc(mem_size, GFP_KERNEL);
  BAILOUT_MEM(mem)

  retval = kernel_read(bprm->file, 0, mem, dos_header->NtHeaderAddress + sizeof(struct s_nt_header) + (sizeof(struct s_section_header) * nt_header->FileHeader.NumberOfSections));
  BAILOUT_READ(retval)

  nt_header = (struct s_nt_header*) (mem + dos_header->NtHeaderAddress);
  section_headers = (struct s_section_header*)((char*)nt_header + sizeof(struct s_nt_header));

  for(i = 0; i < nt_header->FileHeader.NumberOfSections; ++i)
  {
    retval = kernel_read(bprm->file, section_headers[i].PointerToRawData, mem + section_headers[i].VirtualAddress, section_headers[i].SizeOfRawData);
    BAILOUT_READ(retval)
  }



  printk("\nnumber of sections: %i", nt_header->FileHeader.NumberOfSections);
  printk("\noptional header size: 0x%08x", nt_header->FileHeader.SizeOfOptionalHeader);
  printk("\nnumber of dirs: %i", nt_header->OptionalHeader.NumberOfDirectories);
  for(i = 0; i < nt_header->OptionalHeader.NumberOfDirectories; ++i) { printk("\ndata dir: %i - location: 0x%08x - size: 0x%08x", i, nt_header->OptionalHeader.DataDirectory[i].VirtualAddress, nt_header->OptionalHeader.DataDirectory[i].Size); }
  for(i = 0; i < nt_header->FileHeader.NumberOfSections; ++i) { printk("\nsection - offset: 0x%p - name: %.8s - file pointer: 0x%08x", (void*)((void*)(&section_headers[i]) - (void*)mem), section_headers[i].Name, section_headers[i].PointerToRawData); }

  goto load_out;

  setup_new_exec(bprm);
  current;


  set_binfmt(&pe_format);


loader_err:

load_out:

free_mem_out:
  if(mem) { kfree(mem); }
out:
  return retval;
}

static int load_pe_library(struct file* file)
{
  int retval = 0;
  printk("\nchecking library for DLL signature");

  return retval;
}

static int pe_core_dump(long signr, struct pt_regs* regs, struct file* file, unsigned long limit)
{
  int error = 0;
  return error;
}

static int __init init_pe_binfmt(void)
{
  printk(KERN_INFO "Loading Portable Executable linker/loader module.\n");

  return(register_binfmt(&pe_format));
}

static void __exit exit_pe_binfmt(void)
{
  printk(KERN_INFO "Unloading Portable Executable linker/loader module.\n");
  unregister_binfmt(&pe_format);
}

core_initcall(init_pe_binfmt);
module_exit(exit_pe_binfmt);

MODULE_LICENSE("Public Domain");
MODULE_AUTHOR("Daniel Wingerd");
MODULE_DESCRIPTION("enables Portable Executable files (Win32 EXEs) to be run as a native format");